import numpy as np
import random
import cv2

from pathlib import Path
from tensorflow.keras.models import load_model

from train_pac.normal.SimpleDataSetLoader import Simple_DataSet_Loader
from train_pac.normal.SimpleProcessor import Simple_Processor
from img_2_array_preprocess import Img_2_Array_Preprocess


def main():
    print("[info]:加载图像中...")
    base_url = Path("../../dataset/dog_cat")
    pic = [f'{str(base_url)}/cat.{i}.jpg'
           for i in random.choices(np.arange(0, 12500), k=10)]
    pic.extend([f'{str(base_url)}/dog.{i}.jpg'
                for i in random.choices(np.arange(0, 12500), k=10)])
    sp = Simple_Processor(32, 32)
    i2a = Img_2_Array_Preprocess()
    sdl = Simple_DataSet_Loader(preprocessors=[sp, i2a])
    data, labels = sdl.load(pic)
    data = data.astype("float") / 255.0
    print("[info]:加载模型中...")
    model = load_model("shallowNet_animal.hdf5")
    print("[info]:识别图像中...")
    preds = model.predict(data, batch_size=32).argmax(axis=1)
    class_label = ["cat", "dog"]

    for num, temp in enumerate(pic):
        image = cv2.imread(temp)
        cv2.putText(image, f"label:{class_label[preds[num]]}",
                    (10,30), cv2.FONT_HERSHEY_SIMPLEX,0.7, (0,255,0),2)
        cv2.imshow("Image", image)
        cv2.waitKey(0)


if __name__ == '__main__':
    main()
